#!/usr/bin/env python3
import os
import numpy as np
import matplotlib.pyplot as plt
from netCDF4 import Dataset
from wrf import getvar, smooth2d

# ---------- CONFIG ----------
BASE_DIR = "/path/to/base-directory/"
CASES = ["ck_0.3", "ck_0.075", "ck_def"]  # exactly the folder names under BASE_DIR
HURRICANE_NAME = "Sam"
GRID_SIZE = 6  # neighborhood size used in both wind and SLP processing

# ---------- PLOTTING STYLE (from your plotting folder) ----------
plt.rcParams['font.family'] = 'Times New Roman'
plt.rcParams['mathtext.fontset'] = 'custom'
plt.rcParams['mathtext.rm'] = 'Times New Roman'
plt.rcParams['mathtext.it'] = 'Times New Roman:italic'
plt.rcParams['font.size'] = 22

colors = {
    "ck_0.3":  "#FF4500",
    "ck_0.075":"#4682B4",
    "ck_def":  "#808080",
}
markers = {
    "ck_0.3":  "o",
    "ck_0.075":"s",
    "ck_def":  "^",
}

def case_label(key):
    if key == "ck_0.075": return r'c$_k$_0.075'
    if key == "ck_0.3":   return r'c$_k$_0.3'
    if key == "ck_def":   return r'c$_k$_def'
    return key

# ---------- HELPERS ----------
def list_wrfs(d):
    return sorted([f for f in os.listdir(d) if f.startswith("wrfout_d03")])

def window_mean_extrema(arr, k, mode="max"):
    n0, n1 = arr.shape
    best = -np.inf if mode == "max" else np.inf
    hk = k // 2
    for i in range(n0):
        i0, i1 = max(i - hk, 0), min(i + hk + 1, n0)
        for j in range(n1):
            j0, j1 = max(j - hk, 0), min(j + hk + 1, n1)
            m = arr[i0:i1, j0:j1].mean()
            if mode == "max":
                if m > best: best = m
            else:
                if m < best: best = m
    return best

def max_wind_from_file(nc_path, k):
    with Dataset(nc_path, "r") as nc:
        u10 = nc.variables["U10"][0, :, :]
        v10 = nc.variables["V10"][0, :, :]
    ws = np.sqrt(u10**2 + v10**2)
    return window_mean_extrema(ws, k, mode="max")

def min_slp_from_file(nc_path, k):
    with Dataset(nc_path, "r") as nc:
        slp_da = getvar(nc, "slp", timeidx=0)
        slp_sm = smooth2d(slp_da, 3, cenweight=4)
    slp = np.array(slp_sm)
    return window_mean_extrema(slp, k, mode="min")

def process_case_series(case_dir, k):
    files = list_wrfs(case_dir)
    wind_series, slp_series = [], []
    for f in files:
        ncpath = os.path.join(case_dir, f)
        wind_series.append(max_wind_from_file(ncpath, k))
        slp_series.append(min_slp_from_file(ncpath, k))
    return np.array(wind_series), np.array(slp_series)

# ---------- RUN ----------
all_wind = {}
all_slp = {}
for c in CASES:
    cdir = os.path.join(BASE_DIR, c)
    wind_ts, slp_ts = process_case_series(cdir, GRID_SIZE)
    all_wind[c] = wind_ts
    all_slp[c] = slp_ts

# ---------- PLOT (1 hurricane, 3 cases) ----------
fig, axs = plt.subplots(2, 1, figsize=(12, 10))  # (a) intensity, (b) min SLP

# Wind intensity
ax_w = axs[0]
for c in CASES:
    t = np.arange(len(all_wind[c]))
    ax_w.plot(t, all_wind[c], linestyle="--", linewidth=1, marker=markers.get(c, None), color=colors.get(c, None), label=case_label(c))
ax_w.set_title(f"(a) {HURRICANE_NAME} Intensity", fontsize=28)
ax_w.set_ylabel(r'$\mathit{Wind\ Speed}$ (m s$^{-1}$)', fontsize=24)
ax_w.grid(True)
ax_w.tick_params(axis='both', which='major', labelsize=22)

# Minimum SLP
ax_p = axs[1]
for c in CASES:
    t = np.arange(len(all_slp[c]))
    ax_p.plot(t, all_slp[c], linestyle="--", linewidth=1, marker=markers.get(c, None), color=colors.get(c, None), label=case_label(c))
ax_p.set_title(f"(b) {HURRICANE_NAME} Min-SLP", fontsize=28)
ax_p.set_ylabel(r'$\mathit{Minimum\ SLP}$ (hPa)', fontsize=24)
ax_p.set_xlabel(r'Time (hours)', fontsize=24)
ax_p.grid(True)
ax_p.tick_params(axis='both', which='major', labelsize=22)

handles, labels = ax_w.get_legend_handles_labels()
fig.legend(handles, labels, loc='upper center', ncol=len(CASES), bbox_to_anchor=(0.5, 1.02), fontsize=26)
plt.tight_layout(rect=[0, 0, 1, 0.96])
plt.savefig("series.png", dpi=300, bbox_inches="tight")
plt.show()

